# _*_ coding: utf-8 _*_

import numpy as np
from copy import deepcopy
from random import sample
from random import uniform

from scripts.utils.operators import Lifecycle
from scripts.utils.operators import lhs_uniform

np.set_printoptions(linewidth=999, threshold=round(1e9))

class LLSO:
    """
    Level-based Learning Swarm Optimizer (LLSO)
    https://doi.org/10.1109/TEVC.2017.2743016
    """
    def __init__(self, obj_func, dimension: int, bounds: tuple, NP: int, L: int, phi: float):
        self.obj_func = obj_func
        self.__d, self.__bounds, self.__NP, self.__L, self.__phi = dimension, bounds, NP, L, phi

        lhs = lhs_uniform(n=self.__NP, p=self.__d)
        x = (self.__bounds[0] + lhs * (self.__bounds[1] - self.__bounds[0]))  # position
        v = np.zeros(shape=(self.__NP, self.__d))  # velocity
        # population[i] structure: [x, v, fitness, layer]
        self.__population = [[x[i], v[i], 0, 1] for i in range(self.__NP)]
        self.__evaluate()
        self.gBest = deepcopy(min(self.__population, key=lambda _: _[2]))

        self.life = Lifecycle()
        self.life.update(eva=self.__NP)

    def optimize(self, iterations=None, evaluations=None, deadline=None):
        self.life.set_stopping_condition(iterations=iterations, evaluations=evaluations, deadline=deadline)
        while self.life.check_stopping_condition() is False:
            self.__layering()
            self.__learning()
            self.__evaluate()
            if self.gBest[2] > min(self.__population, key=lambda x: x[2])[2]:
                self.gBest = deepcopy(min(self.__population, key=lambda _: _[2]))

            self.life.update(it=1, eva=self.__NP, gBest=self.gBest[2])

        return self.life

    def __evaluate(self):
        for i, x in enumerate(self.__population):
            self.__population[i][2] = self.obj_func(x[0])

    def __layering(self):
        self.__population = sorted(self.__population, key=lambda x: x[2], reverse=False)
        grid, layers = self.__NP / self.__L, 1
        for i, p in enumerate(self.__population):
            self.__population[i][3] = layers
            if (i + 1) == grid * layers:
                layers += 1

    def __learning(self):
        for i in range(self.__NP - 1, -1, -1):  # learning progress starts at the lowest level
            if self.__population[i][3] > 2:
                layer = self.__population[i][3]  # suppose layer >= 3
                k1, k2 = sorted(sample(range(1, layer), k=2))  # randomly select 2 higher levels
                Lk1 = [_ for _ in self.__population if _[3] == k1]  # find Level k1
                Lk2 = [_ for _ in self.__population if _[3] == k2]  # find Level k2
                x1, x2 = sample(Lk1, k=1)[0], sample(Lk2, k=1)[0]  # randomly sample from Lk1, Lk2
                x1, x2 = x1[0], x2[0]  # x1 is better than x2 implicitly (Lk1 < Lk2)
                # update velocity and position
                self.__population[i][1], self.__population[i][0] = self.__update(self.__population[i], x1, x2)

            elif self.__population[i][3] == 2:  # L2
                L1 = [_ for _ in self.__population if _[3] == 1]
                # randomly select 2 particles from L1
                x1, x2 = sorted(sample(L1, k=2), key=lambda x: x[2])  # sorted to make sure x1 is better than x2
                x1, x2 = x1[0], x2[0]  # choose position vector
                # update velocity and position
                self.__population[i][1], self.__population[i][0] = self.__update(self.__population[i], x1, x2)

            elif self.__population[i][3] == 1:  # L1
                continue

    def __update(self, p, x1, x2):
        """
        update velocity and position
        :param p: current particle
        :param x1: learning position 1
        :param x2: learning position 2
        :return: velocity and position
        """
        x, v = p[0], p[1]
        r1 = np.random.uniform(low=0.0, high=1.0, size=self.__d)
        r2 = np.random.uniform(low=0.0, high=1.0, size=self.__d)
        r3 = np.random.uniform(low=0.0, high=1.0, size=self.__d)
        v = r1 * v + r2 * (x1 - x) + self.__phi * r3 * (x2 - x)
        x = x + v
        for i in range(len(x)):
            if x[i] < self.__bounds[0]:
                x[i] = self.__bounds[0]
            if x[i] > self.__bounds[1]:
                x[i] = self.__bounds[1]

        return v, x
